import abc
import json
import os
import pickle
from abc import abstractmethod
from datetime import datetime, timedelta


def get_format_attr(source_name: str, age: tuple[int | None, int | None], gender: int | None, old_attr: str) -> str:
    """
    构造唯一属性名，将 age 区间、gender 和原始属性拼接成统一格式：
    例如：GeoStyle___age18-25___gender1___clothing_category__Dress

    - age: (lower, upper)，None 表示缺失
    - gender: 0=女性, 1=男性, None=未知
    - old_attr: 原始属性名，如 "clothing_category__Dress"
    """
    # 年龄格式化
    if age[0] is None and age[1] is None:
        age_str = "ageNA"
    else:
        lower = str(age[0]) if age[0] is not None else "NA"
        upper = str(age[1]) if age[1] is not None else "NA"
        age_str = f"age{lower}-{upper}"

    # 性别格式化
    gender_map = {0: "F", 1: "M", None: "NA"}
    gender_str = f"gender{gender_map.get(gender, 'NA')}"

    return f"{source_name}___{age_str}___{gender_str}___{old_attr}"


def parse_attr(formatted_attr: str) -> tuple[str, tuple[int | None, int | None], int | None, str]:
    """
    解析格式化的属性字符串，返回 (source_name, age区间, gender, old_attr)

    输入样式:
        "Fit___age18-25___genderM___clothing_category__Pants"
        "GeoStyle___ageNA___genderNA___clothing_category__Dress"
    """
    try:
        source_name, age_part, gender_part, old_attr = formatted_attr.split("___", 3)

        # 解析年龄
        if age_part == "ageNA":
            age = (None, None)
        else:
            parts = age_part[3:].split("-")
            age = (
                int(parts[0]) if parts[0] != "NA" else None,
                int(parts[1]) if parts[1] != "NA" else None
            )

        # 解析性别
        gender_map = {"F": 0, "M": 1, "NA": None}
        gender_str = gender_part[6:]
        gender = gender_map.get(gender_str, None)

        return source_name, age, gender, old_attr

    except Exception as e:
        raise ValueError(f"Invalid formatted attribute string: {formatted_attr}") from e


class BaseDataset(metaclass=abc.ABCMeta):
    source_name = None
    start_time = None
    end_time = None
    delta = None

    @abstractmethod
    def get_city_name_list(self):
        pass

    @abstractmethod
    def get_attributes_list(self, target_city_name: str):
        pass

    @abstractmethod
    def get_preprocessed_normed(self):
        pass

    @abstractmethod
    def get_preprocessed_norm(self):
        pass


class GeoStyleDataset(BaseDataset):
    """
    2013 - 2015
    每周一个点
    """

    source_name = "GeoStyle"
    start_time = '2013-08-05'
    end_time = '2016-05-02'
    delta = 7

    def __init__(self, data_dir: str):
        self._data_dir = data_dir
        # {城市ID: {服装属性: [[时间戳, 原始趋势值], ...]}}
        raw_data_path = os.path.join(self._data_dir, "geostyle_data.json")
        with open(raw_data_path, 'r') as fs:
            self._raw = json.load(fs)

        # {城市ID: {服装属性: [最小值, 最大值, 平均增量]}}
        norm_data_path = os.path.join(self._data_dir, "geostyle_data_norm.json")
        with open(norm_data_path, 'r') as fs:
            self._norm = json.load(fs)

        # {城市ID: {服装属性: [[时间戳, 已归一化趋势值], ...]}}
        normed_data_path = os.path.join(self._data_dir, "geostyle_data_normed.json")
        with open(normed_data_path, 'r') as fs:
            self._normed = json.load(fs)

        metadata_path = os.path.join(self._data_dir, "metadata.pkl")
        with open(metadata_path, 'rb') as fs:
            self._metadata = pickle.load(fs)

        def load_city_id_mapping_name_dict():
            """
            获取数据集城市列表
            """
            return self._metadata['cities']

        def process_normed():
            """
            key: (city_name, attr)
            """
            result = {}
            for city_id, city_data in self._normed.items():
                city_name = self._city_id_mapping_name_dict[city_id]

                for attr, ts_list in city_data.items():
                    attr = get_format_attr(self.source_name, (None, None), None, attr)
                    key = (city_name, attr)

                    processed_ts = []
                    date = datetime(2013, 8, 5)

                    for week_id, value in ts_list:
                        processed_ts.append((date, value))
                        date = date + timedelta(days=GeoStyleDataset.delta)

                    result[key] = processed_ts

            return result

        def load_attributes():
            """
            加载每个城市中的属性
            """
            res = []
            for i in range(len(self._metadata['attributes'])):
                for j in range(len(self._metadata['categories'][i])):
                    old_attr = f"{self._metadata['attributes'][i]}__{self._metadata['categories'][i][j]}"
                    attr = get_format_attr(self.source_name, (None, None), None, old_attr)
                    res.append(attr)
            return res

        def process_norm():
            result = {}
            for city_id, city_data in self._norm.items():
                city_name = self._city_id_mapping_name_dict[city_id]

                for old_attr, norm_list in city_data.items():
                    attr = get_format_attr(self.source_name, (None, None), None, old_attr)
                    key = (city_name, attr)
                    result[key] = norm_list
            return result

        self._city_id_mapping_name_dict = load_city_id_mapping_name_dict()
        self._city_name_list = list(self._city_id_mapping_name_dict.values())
        self._preprocessed_normed = process_normed()
        self._attributes_list = load_attributes()
        self._preprocessed_norm = process_norm()

    def get_city_name_list(self):
        """
        获取geostyle数据集城市列表
        """
        return self._city_name_list

    def get_attributes_list(self, target_city_name=None):
        """
        获取属性列表
        """
        return self._attributes_list

    def get_preprocessed_normed(self):
        return self._preprocessed_normed

    def get_preprocessed_norm(self):
        return self._preprocessed_norm


class FITDataset(BaseDataset):
    """
    2014.7 - 2019-6
    半个月一个点
    """
    source_name = "FIT"
    start_time = '2014-07-15'
    end_time = '2019-06-04'
    delta = 15

    def __init__(self, data_dir: str):
        self._data_dir = data_dir

        # 加载数据
        with open(os.path.join(self._data_dir, "FIT_data.json"), "r") as fs:
            self._raw = json.load(fs)
        # 加载标准化的数据
        with open(os.path.join(self._data_dir, "FIT_data_norm.json"), "r") as fs:
            self._norm = json.load(fs)
        # 加载知识图谱邻接矩阵
        with open(os.path.join(self._data_dir, "kg_adj_normed.json"), "r") as fs:
            self._kg_adj = json.load(fs)
        # 加载知识图谱组邻接矩阵
        with open(os.path.join(self._data_dir, "kg_group_adj_normed.json"), "r") as fs:
            self._kg_group_adj = json.load(fs)

        def process_fit():
            """
            处理 FIT 数据集，输出格式:
            key: (city_name, formatted_attr)
            value: List[[datetime, value]]
            """
            result = {}

            for all_attr_key, attr_dict in self._raw.items():
                city_name = None
                age = (None, None)
                gender = None
                one_attr = all_attr_key.split('__')

                # 提取 city, age, gender
                for j in one_attr:
                    if j.startswith('city:'):
                        city_name = j.split(':', 1)[1]

                    elif j.startswith('age:'):
                        age_str = j.split(':', 1)[1]
                        if age_str == 'age_lt_18':
                            age = (0, 17)
                        elif age_str == 'age_18_25':
                            age = (18, 25)
                        elif age_str == 'age_25_40':
                            age = (25, 40)
                        elif age_str == 'age_gt_40':
                            age = (41, None)
                        else:
                            age = (None, None)

                    elif j.startswith('gender:'):
                        gender_str = j.split(':', 1)[1].lower()
                        if gender_str == 'male':
                            gender = 1
                        elif gender_str == 'female':
                            gender = 0

                # 遍历属性趋势时间序列
                for old_attr, ts_list in attr_dict.items():
                    formatted_attr = get_format_attr(self.source_name, age, gender, old_attr)
                    key = (city_name, formatted_attr)

                    processed_ts = []
                    base_date = datetime(2014, 7, 15)

                    for period_id, value in ts_list:
                        processed_ts.append((base_date, value))
                        base_date = base_date + timedelta(days=15)

                    result[key] = processed_ts

            return result

        def load_city_name_list():
            """
            获取fit数据集城市列表
            """
            s = set()
            for (city_name, _), v in self._preprocessed_normed.items():
                s.add(city_name)
            return list(s)

        def load_attributes():
            """
            加载每个城市中的属性
            返回dict
            """
            res = {'all': set()}
            for (city_name, attr), v in self._preprocessed_normed.items():
                if res.get(city_name) is None:
                    res[city_name] = set()

                res['all'].add(attr)
                res[city_name].add(attr)

            for k, v in res.items():
                res[k] = list(v)
            return res

        def process_norm():
            result = {}
            for all_attr_key, attr_dict in self._norm.items():
                city_name = None
                age = (None, None)
                gender = None
                one_attr = all_attr_key.split('__')

                # 提取 city, age, gender
                for j in one_attr:
                    if j.startswith('city:'):
                        city_name = j.split(':', 1)[1]

                    elif j.startswith('age:'):
                        age_str = j.split(':', 1)[1]
                        if age_str == 'age_lt_18':
                            age = (0, 17)
                        elif age_str == 'age_18_25':
                            age = (18, 25)
                        elif age_str == 'age_25_40':
                            age = (25, 40)
                        elif age_str == 'age_gt_40':
                            age = (41, None)
                        else:
                            age = (None, None)

                    elif j.startswith('gender:'):
                        gender_str = j.split(':', 1)[1].lower()
                        if gender_str == 'male':
                            gender = 1
                        elif gender_str == 'female':
                            gender = 0
                for old_attr, norm_list in attr_dict.items():
                    attr = get_format_attr(self.source_name, age, gender, old_attr)
                    key = (city_name, attr)
                    result[key] = norm_list
            return result

        self._preprocessed_normed = process_fit()
        self._city_name_list = load_city_name_list()
        self._attributes_dict = load_attributes()
        self._preprocessed_norm = process_norm()

    def get_city_name_list(self):
        """
        获取fit数据集城市列表
        """
        return self._city_name_list

    def get_attributes_list(self, city_name='all'):
        """
        获取属性列表
        如果city_name != None，返回特定城市中的属性
        """
        return self._attributes_dict.get(city_name, [])

    def get_preprocessed_normed(self):
        return self._preprocessed_normed

    def get_preprocessed_norm(self):
        return self._preprocessed_norm
